from chartspy import Echarts, KlineCharts
from chartspy import express
import pandas as pd
import datetime as dt

from shares.shorts import tsc


def kline_charts(ts_code, start_date, end_date, adj="qfq", main_indicators=['MA'],
                 mas=[5, 10, 30, 60, 120, 250])->KlineCharts:
    """
    展示带常用指标K线
    """
    start_date = (dt.datetime.now() - dt.timedelta(days=1000)).strftime(
        "%Y-%m-%d") if start_date is None else start_date
    end_date = dt.datetime.now().strftime("%Y-%m-%d") if end_date is None else end_date
    df_price = tsc.pro_bar(ts_code=ts_code, start_date=start_date, end_date=end_date, adj=adj).sort_values("trade_date")
    df_price = df_price.rename(columns={'trade_date': 'timestamp', 'vol': 'volume'})
    return KlineCharts(df_price, main_indicators=main_indicators, mas=mas)


def candlestick(ts_code, start_date, end_date, adj="qfq", mas=[5, 10, 30, 60, 120, 250], height='600px')->Echarts:
    """
    展示K线
    """
    start_date = (dt.datetime.now() - dt.timedelta(days=1000)).strftime(
        "%Y-%m-%d") if start_date is None else start_date
    end_date = dt.datetime.now().strftime("%Y-%m-%d") if end_date is None else end_date
    df_price = tsc.pro_bar(ts_code=ts_code, start_date=start_date, end_date=end_date, adj=adj).sort_values("trade_date")
    return express.candlestick_echarts(df_price, time_field='trade_date', volume_field='vol', mas=mas, height=height)


def drawdown(ts_code, start_date, end_date, adj="qfq", height='600px')->Echarts:
    start_date = (dt.datetime.now() - dt.timedelta(days=1000)).strftime(
        "%Y-%m-%d") if start_date is None else start_date
    end_date = dt.datetime.now().strftime("%Y-%m-%d") if end_date is None else end_date
    df_price = tsc.pro_bar(ts_code=ts_code, start_date=start_date, end_date=end_date, adj=adj).sort_values("trade_date")
    return express.drawdown_echarts(df_price, time_field='trade_date', value_field='close', code_field='ts_code',
                                    height=height)


def show_signal(ts_code, date, title="信号", left_extend_days=255, right_extend_days=255, height="600px")->Echarts:
    """
    显示单一信号
    """
    day = pd.to_datetime(date)
    start_date = (day - dt.timedelta(days=left_extend_days)).strftime("%Y-%m-%d")
    end_date = (day + dt.timedelta(days=right_extend_days)).strftime("%Y-%m-%d")
    df_price = tsc.pro_bar(ts_code=ts_code, start_date=start_date, end_date=end_date, adj='qfq').sort_values(
        "trade_date")
    df_segment = pd.DataFrame(columns=['date', 'label'], data=[[day.strftime("%Y-%m-%d"), "信号日"]])
    return express.candlestick_echarts(df_price, volume_field='vol', time_field='trade_date', log_y=True,
                                       title=title, height=height).overlap_series(
        [express.mark_vertical_line_echarts(df_segment, x='date', label='label', title='信号')])


def show_signals(ts_code, df_signal, date="date", category='category', title="信号", left_extend_days=255,
                 right_extend_days=255, height="600px")->Echarts:
    """
    显示信号集合
    df_signal [date,category]
    """
    signals = df_signal.copy()
    signals[date] = pd.to_datetime(signals[date])
    max_date = (pd.to_datetime(signals[date].max()) + dt.timedelta(days=right_extend_days)).strftime("%Y-%m-%d")
    min_date = (pd.to_datetime(signals[date].min()) - dt.timedelta(days=right_extend_days)).strftime("%Y-%m-%d")
    df_price = tsc.pro_bar(ts_code=ts_code, start_date=min_date, end_date=max_date, adj='qfq').sort_values(
        "trade_date")
    df_price['trade_date'] = pd.to_datetime(df_price['trade_date'])
    signals = signals.merge(df_price, left_on=[date], right_on=['trade_date'])
    return express.candlestick_echarts(df_price, volume_field='vol', time_field='trade_date', log_y=True,
                                       title=title, height=height).overlap_series(
        [express.mark_vertical_line_echarts(signals, x=date, label=category, title="信号")
         ]
    )
